// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*! 
\file
\brief Main Body of PASSOC_BT_FT2 program (Boolean Data)
*/

/*! 
\defgroup	passoc_bt_ft2 PASSOC_BT_FT2 - Basic Association test (BT)
\ingroup    mpi_stats
\brief		PASSOC_BT_FT2 - Basic Association test (BT)
\details
<p>
This program performs Basic Association tests between SNPs and a binary trait.
This program stores allelic data as boolean arrays instead of BitString as used by the other programs.
Although this does mean that it uses a lot of memory when parsing large data-sets,
performance tests obtained 15% increase in performance when compared to a version of the program
which used BitString for storage of allelic data.
</p>

<p>This program runs in 2 modes</p>
<dl>
<dt>Mode 1</dt>
<dd>Multi-core MPI invocations which perform statistical calculation with permutation.</dd>
<dt>Mode 2</dt>
<dd>Single-core MPI call which parses the permutation data and creates the final output file.</dd>
</dl>

<p>By default this program runs in <em>Mode 1</em>, however if the <em>--merge</em> parameter is specified, the 
program runs in the <em>Mode 2</em> where it attempts to read the permutation data.<br>
If multiple cores are assigned to a run when <em>--merge</em> is specified, PASSOC_BT_FT2 exits with an error code.
</p>
<p>The statistical output is placed in a file with the suffix \ref SUFFIX_ASSOC_FILE .</p>
<p>If successful PASSOC_BT_FT returns with an exitcode of 0, if it encountered an error the programs returns 
with an exit code of -1.</p>

<h2>OPTIONS</H2>

<table class=options border=0>
<tr><th class=param_header colspan=1>PARAMETER</th>
<th class=param_header colspan=3>DESCRIPTION</th></tr>

<tr><td class=param_name colspan=1>--id STRING</td>
<td class=param_desc colspan=3>Unique identifer for a cluster job (also sets output file stem).</td></tr>

<tr><td class=param_name colspan=1>--rdir PATH</td>
<td class=param_desc colspan=3>Path to remote working directory on central file system.</td></tr>

<tr><td class=param_name colspan=1>--ldir PATH</td>
<td class=param_desc colspan=3>Path to local working directory on each core of the cluster.</td></tr>

<tr><td class=param_name colspan=1>--bmap FILE</td>
<td class=param_desc colspan=3>Path to the BMAP file (binary-encoded SNP list).</td></tr>

<tr><td class=param_name colspan=1>--samples FILE</td>
<td class=param_desc colspan=3>Path to the Sample file (people and phenotype).</td></tr>

<tr><td class=param_name colspan=1>--sped FILE</td>
<td class=param_desc colspan=3>Path to the SPED file (Allelic data, SNP Major, bit-strings).</td></tr>

<tr><td class=param_name colspan=1>--perm FILE</td>
<td class=param_desc colspan=3>Path to the permutation indices file.</td></tr>

<tr><td class=param_name colspan=1>--merge</td>
<td class=param_desc colspan=3>Signal PASSOC_BT_FT2 to merge permutation results and create final output file (single-core operation).</td></tr>

<tr><td class=param_name colspan=1>--missed-rank 0..101</td>
<td class=param_desc colspan=3>%Set processor rank to number specified by this parameter instead of rank assigned by MPI.<br>
To be used only if a section of a calculation needs to run a single core if a processor failed.
</td></tr>

<tr><td class=param_name colspan=1>--missed-procs 1..102</td>
<td class=param_desc colspan=3>%Set the No. processors associated with a calculation incase processor rank set by <em>missed-rank</em>.</td></tr>

<tr><td class=param_name colspan=1>--ci 0..1.0</td>
<td class=param_desc colspan=3>%Set the confidence interval (CI) of the Odds-Ratio (OR) for the dominant, recessive, trend and allelic tests.<br>
Default ci=0.95
</td></tr>

<tr><td class=param_name colspan=1>--sort-by COLUMN_NAME</td>
<td class=param_desc colspan=3>Sort column for the final CSV file containing the statistical output.<br>
Default sort column=EMP1.
</td></tr>
</table>
*/

#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <mpi.h>

#include "SMG_All.h"

/*!
\brief Command-line options of PASSOC_BT_FT2
*/
struct option long_options[] = {
	{ "bmap", required_argument, NULL, 'm' }, // Path to the map file.
	{ "samples", required_argument, NULL, 's' }, // Path to the map file.
	{ "sped", required_argument, NULL, 'p' }, // Path to the PED file.
	{ "id", required_argument, NULL, 'i' }, // Run ID
	{ "rdir", required_argument, NULL, 'r'}, // Remote Output direct
	{ "ldir", required_argument, NULL, 'l'}, // Local Output direct
	{ "merge", optional_argument, NULL, 'M'}, // Flag about to do the merge.
	{ "missed-rank", required_argument, NULL, 'n'}, // A rank that failed for whatever reason that needs re-running.
	{ "missed-procs", required_argument, NULL, 'o'}, // No. processors associated with the original calculation.
	{ "ci", required_argument, NULL, 'c' }, // Confidence interval
	{ "sort-by", required_argument, NULL, 'S' }, // Column Name to sort final output by.
	{ "nperm", required_argument, NULL, 'w'}, // No. permutations.
	{ "perm", required_argument, NULL, 'P'}, // Path to permutation file (superfluous parameter now).
	{ 0,    0,    0,    0   } 
};

/*!
\brief Main body of the PASSOC_BT_FT2 program
\ingroup passoc_bt_ft2
*/
int passoc_bt2_main(int argc, char *argv[]) {
	char ch, *spedData = NULL, *data = NULL, *bs = NULL;
	struct selected_options ops;
	int myrank, numprocs, result;
	BOOL jobDone = FALSE;
	struct chromosomes chromos;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];
	int nArgCount = 0, nRequiredArgs = 6, i, j;
	struct small_locus *map = NULL;
	struct sample *samples = NULL;
	int nMap = 0, nSamples = 0, nSpedSamples = 0, nSNPs = 0, nSpedRecordSizeInBytes = 0;
	BOOL isDoingMerge = FALSE;
	double *originals = NULL, *results = NULL, zt;
	struct adjusted *adj = NULL;
	int *R = NULL, *maxR = NULL, missedRank = -1, missedProcs = -1, nAdj = 0, *perms = NULL, *seeds = NULL, nSeeds = 0;

	// Initialise MPI environment, get size and rank.
	if ((result = MPI_Init(&argc,&argv)) != MPI_SUCCESS) mpi_error(result, "Initialisation");
	if ((result = MPI_Comm_size(MPI_COMM_WORLD,&numprocs)) != MPI_SUCCESS) mpi_error(result, "Number of processors");
	if ((result = MPI_Comm_rank(MPI_COMM_WORLD,&myrank)) != MPI_SUCCESS) mpi_error(result, "Processor number");

	init_genrand((unsigned long) time(NULL));
	options_init(&ops);
	chromosomes_init(&chromos);
	chromosomes_define_human(&chromos);

	// Want to set these set via the command-line.
	*ops.szRunId = *ops.ldir = *ops.rdir = '\0'; 
	ops.verbose = FALSE;

	while ((ch = (char) getopt_long(argc, argv, "m:v:p:f:i:", long_options, NULL)) != -1) {
		double value = 0.0;
		switch (ch) {
			case 'i':
				strcpy(ops.szRunId, optarg);
				nArgCount++;
				break;
			case 's': // Sample File
				strcpy(ops.szSampleFileName, optarg);
				nArgCount++;
				break;
			case 'p': // SPED file
				strcpy(ops.szSpedFileName, optarg);
				nArgCount++;
				break;
			case 'm': // Bmap File
				strcpy(ops.szBmapFileName, optarg);
				nArgCount++;
				break;
			case 'r': // Remote output directory
				strcpy(ops.rdir, optarg);
				nArgCount++;
				break;
			case 'l': // Local output directory
				strcpy(ops.ldir, optarg);
				nArgCount++;
				break;
			case 'M':
				isDoingMerge = TRUE;
				break;
			case 'n': // Missed rank to re-run calculation on a single processor.
				missedRank = atoi(optarg);
				break;
			case 'o': // No. processors associated with a 'missed' rank.
				missedProcs = atoi(optarg);
				break;
			case 'c': // The confidence interval
				value = atof(optarg);
				if (value >= 0.5 && value <= 1.0) {
					ops.ci_level = value;
				}
				break;
			case 'S': // Sort column name
				if (strlen(optarg) > 0) strcpy(ops.szSortColumnName, optarg);
				break;
			case 'w': // Number of Permutations.
				ops.nperm = atoi(optarg);
				if (ops.nperm == 0) ops.nperm = 10000;
				break;
			default:
				break;
		}
	}
	zt = ltqnorm(1 - (1 - ops.ci_level) / 2); // For the CI of the OR

	// If the required arguements are not specified, there's no point carrying on.
	if (isDoingMerge) nRequiredArgs = 3; // just need ID, ops.ldir & ops.rdir
	if (nArgCount != nRequiredArgs) {
		chromosomes_free_attributes(&chromos);
		MPI_Finalize();
		return -1;
	}

	// See if a calculation needs to be re-run on a 'single' processor.
	// Check that all the command-line options are set.
	if (missedRank >= 0 || missedProcs >= 0) {
		BOOL exitEarly = FALSE;

		if (numprocs != 1) exitEarly = TRUE;
		else if (missedRank == -1 || missedProcs == -1) exitEarly = TRUE;
		else if (!(missedRank >= 0 && missedRank <= missedProcs)) exitEarly = TRUE;

		if (exitEarly) {
			MPI_Finalize();
			return -1;
		} else {
			myrank = missedRank;
			numprocs = missedProcs;
		}
	} 

	if (!isDoingMerge) {
		// Do the non-permuted association calculation across all ranks 
		// in the cluster allocated to this job.
		if ((map = bmapfile_load(&ops, &nMap)) != NULL && 
			(samples = samplefile_load(&ops, &nSamples)) != NULL &&
			(spedData = spedfile_load(&ops, &nSNPs, &nSpedSamples, &nSpedRecordSizeInBytes)) != NULL &&
			(seeds = seedfile_load(&ops, &nSeeds)) != NULL)
		{
			data = bitstring_unpack(spedData, nSNPs, nSpedSamples, nSpedRecordSizeInBytes);
			if (data != NULL) {
				free(spedData);
				spedData = NULL;
			}
			if (nMap > 0 && nSNPs == nMap && data != NULL) {
				BOOL adjAllocated = TRUE;
				int nPermsPerProcessor = get_tasks_per_processor(numprocs, ops.nperm);
				int start_perm, end_perm, i;

				start_perm = myrank * nPermsPerProcessor;
				end_perm = start_perm + nPermsPerProcessor;
				if (end_perm > ops.nperm || myrank == numprocs - 1) end_perm = ops.nperm;

				originals = create_vector(nSNPs, 0);
				results = create_vector(nSNPs, 0);
				R = create_int_array(nSNPs, 0);
				perms = create_int_array(nSamples, 0); // Permutation indices.
				maxR = create_int_array(nSNPs, 0);
				if (myrank == ROOT) {
					adj = adjusted_create_array(nSNPs);
					if (adj == NULL) adjAllocated = FALSE;
				}

				// Permute the data-set.
				if (originals != NULL && R != NULL && maxR != NULL && results != NULL && adjAllocated && perms != NULL) {					
					samplefile_aff_coding(samples, nSamples);
					samplefile_set_original_order(samples, nSamples);
					init_genrand(seeds[myrank]);

					// Non-permuted calculation.
					for (i = 0, bs = data; i < nSNPs; i++, bs += (2 * nSamples)) {
						BOOL X = FALSE, haploid = FALSE;

						// Check if the SNP is located on autosome.
						if (chromos.chr_sex[map[i].chr]) X = TRUE;
						else if (chromos.chr_haploid[map[i].chr]) haploid = TRUE;

						// Calculate the original pvalue and affected counts.
						if (myrank != ROOT) originals[i] = passoc_bt2_assoc(X, haploid, bs, samples, nSamples, NULL, NULL, 0, 0, 0, NULL);
						else originals[i] = passoc_bt2_assoc(X, haploid, bs, samples, nSamples, NULL, NULL, 0, 0, 0, &adj[i]);
					}

					for (i = start_perm; i < end_perm; i++) {	
#ifdef ECHO_PERMUTATION
					if (myrank == ROOT) {
						fprintf(stdout, "Permutation: %i\r", i + 1);
						fflush(stdout);
					}
#endif
						zero_vector(results, nSNPs);
						permute(perms, nSamples);
						samplefile_permute_phenotype(samples, perms, nSamples);
						for (j = 0, bs = data; j < nSNPs; j++, bs += (2 * nSamples)) {
							BOOL X = FALSE, haploid = FALSE;

							// Check if the SNP is located on autosome.
							if (chromos.chr_sex[map[j].chr]) X = TRUE;
							else if (chromos.chr_haploid[map[j].chr]) haploid = TRUE;

							// Calculate the original pvalue and affected counts.
							results[j] = passoc_bt2_assoc(X, haploid, bs, samples, nSamples, NULL, NULL, 0, 0, zt, NULL);
						}
						perm_update(nSNPs, results, originals, R, maxR);
					}

					// Dump the local counts to file.
					jobDone = _count_file_write_and_copy(&ops, SUFFIX_ABC_FILE, myrank, numprocs, abs(end_perm - start_perm), nSNPs, R, maxR) ? TRUE : FALSE;
					if (myrank == ROOT && jobDone) jobDone = adjusted_write_and_copy(adj, nSNPs, &ops, myrank) ? TRUE : FALSE;
				}
			}
		}
	} else {
		// Merging & sorting the output files.
		if (numprocs == 1) { // This code in this mode needs to run as a singleton.

			BOOL *file_copy_failures = NULL, hasCopiedAll = TRUE;
			int nPermutations = 0;

			if (strlen(ops.szBmapFileName) == 0) sprintf(ops.szBmapFileName, "%s%s%s", ops.rdir, ops.szRunId, SUFFIX_BMAP_FILE);
			if (strlen(ops.szSpedFileName) == 0) sprintf(ops.szSpedFileName, "%s%s%s", ops.rdir, ops.szRunId, SUFFIX_SPED_FILE);
			if (strlen(ops.szSampleFileName) == 0) sprintf(ops.szSampleFileName, "%s%s%s", ops.rdir, ops.szRunId, SUFFIX_SAMPLE_FILE);

			map = bmapfile_load(&ops, &nMap);
			R = create_int_array(nMap, 0);
			maxR = create_int_array(nMap, 0);
			file_copy_failures = _count_file_parse(&ops, SUFFIX_ABC_FILE, nMap, &nPermutations, R, maxR, &hasCopiedAll, &numprocs);

			spedData = spedfile_load(&ops, &nSNPs, &nSpedSamples, &nSpedRecordSizeInBytes);
			if (spedData != NULL) {
				data = bitstring_unpack(spedData, nSNPs, nSpedSamples, nSpedRecordSizeInBytes);
				free(spedData);
				spedData = NULL;
			}
			samples = samplefile_load(&ops, &nSamples);
			adj = _adjusted_load(&ops, &nAdj);
			if (map != NULL && nMap > 0 && data != NULL && nSNPs > 0 && nSNPs == nMap && samples != NULL && nSamples > 0 && adj != NULL && nAdj == nMap) {
				if (hasCopiedAll) {
					FILE *out = NULL;

					sprintf(localFilepath, "%s%s%s", ops.ldir, ops.szRunId, SUFFIX_ASSOC_FILE);
					sprintf(remoteFilepath, "%s%s%s", ops.rdir, ops.szRunId, SUFFIX_ASSOC_FILE);
					out = fopen(localFilepath, "w");
					if (out != NULL) {
						struct adjusted *a = NULL; 
						double nperms = (double) nPermutations + 1;
						samplefile_aff_coding(samples, nSamples);
						samplefile_set_original_order(samples, nSamples);

						// Write out the results of each SNP to file.
						passoc_bt_write_file_header2(out, &ops);
						for (i = 0, bs = data, a = adj; i < nSNPs; i++, bs += (2 * nSamples), a++) {
							double emp1, emp2;
							BOOL X = FALSE, haploid = FALSE;

							// Check if the SNP is located on autosome.
							if (chromos.chr_sex[map[i].chr]) X = TRUE;
							else if (chromos.chr_haploid[map[i].chr]) haploid = TRUE;

							// Ratio counts
							emp1 = (double) (R[i] + 1) / nperms;
							emp2 = (double) (maxR[i] + 1) / nperms;

							passoc_bt2_assoc(X, haploid, bs, samples, nSamples, &map[i], out, emp1, emp2, zt, a);
						}
						fclose(out);

						// Final sort and clean up.
						jobDone = sort_file_by_col_and_copy(localFilepath, remoteFilepath, ops.szSortColumnName) ? TRUE : FALSE;	
						if (jobDone) assoc_cleanup_temp_files(&ops, SUFFIX_ABC_FILE, numprocs);
					} else {
						write_failed_rank_log(&ops, numprocs, file_copy_failures);
					}
				}	
			} 
		}
	}

	// Free now as not needed anymore
	if (samples != NULL) free(samples);
	if (spedData != NULL) free(spedData);
	if (map != NULL) free(map); 
	if (originals != NULL) free(originals);
	if (R != NULL) free(R);
	if (maxR != NULL) free(maxR);
	if (results != NULL) free(results);
	if (data != NULL) free(data);
	if (adj != NULL) free(adj);
	if (perms != NULL) free(perms);
	chromosomes_free_attributes(&chromos); 

	MPI_Finalize();

	if (jobDone) return 0;
	else return -1;
}

int main(int argc, char *argv[]) {
	return passoc_bt2_main(argc, argv);
}